import math
import torch
from torch import nn
import torch.nn.functional as F
device = "cuda" if torch.cuda.is_available() else "cpu"

"""基于位置的前馈网络"""
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))

""" add & norm 模块"""
class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)

"""  缩放点积注意力机制模块  """
class DotProductAttention(nn.Module):
    """缩放点积注意力"""

    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        # 初始化dropout层，用于注意力权重的正则化
        self.dropout = nn.Dropout(dropout)
    
    """
       前向传播

       参数:
           query: [batch_size, seq_len_q, d_k] - 查询向量，表示要查询的信息
           key: [batch_size, seq_len_k, d_k] - 键向量，表示被查询的信息  
           value: [batch_size, seq_len_v, d_v] - 值向量，包含实际的信息内容
           mask: [batch_size, seq_len_q, seq_len_k] 或 [seq_len_q, seq_len_k] - 可选掩码，用于屏蔽无效位置

       返回:
           output: [batch_size, seq_len_q, d_v] - 注意力加权后的输出
           attention_weights: [batch_size, seq_len_q, seq_len_k] - 注意力权重矩阵
    """
    def forward(self, query, key, value, mask=None):
        # 获取查询向量维度
        d_k = query.size(-1)
        
        # 计算缩放点积分数
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
        
        # 应用掩码（可选）
        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(0)  # 2D转3D
            elif mask.dim() == 3 and mask.size(0) == 1:
                mask = mask.expand(scores.size(0), -1, -1)  # 扩展到批次
            scores = scores.masked_fill(mask == 0, -1e9)  # 掩码位置设为极小值
        
        # 计算注意力权重
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # 步骤5: 将注意力权重应用于值向量
        # attention_weights: [batch_size, seq_len_q, seq_len_k]
        # value: [batch_size, seq_len_v, d_v] (通常seq_len_v = seq_len_k)
        # 矩阵乘法结果: [batch_size, seq_len_q, d_v]
        output = torch.matmul(attention_weights, value)
        
        return output, attention_weights

"""位置编码模块"""
class PositionalEncoding(nn.Module):
    """
    位置编码模块
    为输入序列添加位置信息，使用正弦和余弦函数生成位置编码
    
    实现原理:
    - 使用正弦和余弦函数的不同频率来编码位置信息
    - 偶数位置使用正弦函数: PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
    - 奇数位置使用余弦函数: PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    - 这种编码方式允许模型学习相对位置关系
    
    关键特性:
    - 支持任意长度的序列（最大长度由max_len参数限制）
    - 位置编码不参与梯度更新（注册为缓冲区）
    - 自动处理不同输入维度格式
    """
    
    def __init__(self, d_model, max_len=5000, dropout=0.1, **kwargs):
        super(PositionalEncoding, self).__init__(**kwargs)
        self.dropout = nn.Dropout(p=dropout)
        
        # 创建位置编码矩阵
        pe = torch.zeros(max_len, d_model) # [max_len, d_model] - 存储所有可能位置的所有维度的编码
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # [max_len, 1]
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # [d_model//2]
        
        # 应用正弦和余弦函数
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        # unsqueeze(0): [1, max_len, d_model] - 添加批次维度
        # transpose(0, 1): [max_len, 1, d_model] - 交换维度以匹配输入格式
        pe = pe.unsqueeze(0).transpose(0, 1)
        
        # 注册为缓冲区（不参与梯度更新）
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """
        前向传播
        
        参数:
            x: [seq_len, batch_size, d_model] 或 [batch_size, seq_len, d_model]
        
        返回:
            x + positional_encoding: 添加位置编码后的输出
        """
        # 维度检查: 如果是批次优先格式 [batch_size, seq_len, d_model]
        if x.dim() == 3 and x.size(1) == self.pe.size(0):
            # 维度变换: [max_len, 1, d_model] → [seq_len, 1, d_model] → [1, seq_len, d_model] → 广播到 [batch_size, seq_len, d_model]
            x = x + self.pe[:x.size(1), :].transpose(0, 1)
        else:  # 序列优先格式 [seq_len, batch_size, d_model]
            # 维度变换: [max_len, 1, d_model] → [seq_len, 1, d_model] → 广播到 [seq_len, batch_size, d_model]
            x = x + self.pe[:x.size(0), :]
        
        return self.dropout(x)

"""编码器层模块"""
class EncoderLayer(nn.Module):
    """
    编码器层
    包含多头自注意力和前馈网络，以及残差连接和层归一化
    """
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, **kwargs):
        super(EncoderLayer, self).__init__(**kwargs)
        
        # 多头自注意力机制
        self.self_attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # 前馈网络
        self.feed_forward = PositionWiseFFN(d_model, d_ff, d_model)
        
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # 残差连接和dropout
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        """
        前向传播
        
        参数:
            x: [batch_size, seq_len, d_model] - 输入序列
            mask: [batch_size, seq_len, seq_len] - 可选掩码
        
        返回:
            output: [batch_size, seq_len, d_model] - 编码后的输出
        """
        # 多头自注意力 + 残差连接 + 层归一化
        attn_output, _ = self.self_attention(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))
        
        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout2(ff_output))
        
        return x

"""编码器模块"""
class Encoder(nn.Module):
    """
    编码器
    由多个编码器层堆叠而成，包含位置编码和层归一化
    """
    
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, 
                 max_len=5000, dropout=0.1, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        
        self.d_model = d_model
        
        # 词嵌入层
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # 位置编码
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 编码器层堆叠
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        # 输出层归一化
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        """
        前向传播
        
        参数:
            x: [batch_size, seq_len] - 输入序列的索引
            mask: [batch_size, seq_len, seq_len] - 可选掩码
        
        返回:
            output: [batch_size, seq_len, d_model] - 编码后的输出
        """
        # 词嵌入 + 缩放
        x = self.embedding(x) * math.sqrt(self.d_model)
        
        # 位置编码
        x = self.pos_encoding(x)
        
        # 通过所有编码器层
        for layer in self.layers:
            x = layer(x, mask)
        
        # 最终层归一化
        x = self.norm(x)
        
        return x

"""解码器层模块"""
class DecoderLayer(nn.Module):
    """
    解码器层
    包含掩码多头自注意力、编码器-解码器注意力和前馈网络
    """
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, **kwargs):
        super(DecoderLayer, self).__init__(**kwargs)
        
        # 掩码多头自注意力（自回归注意力）
        self.self_attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # 编码器-解码器注意力
        self.cross_attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # 前馈网络
        self.feed_forward = PositionWiseFFN(d_model, d_ff, d_model)
        
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        
        # 残差连接和dropout
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
    
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        """
        前向传播
        
        参数:
            x: [batch_size, tgt_seq_len, d_model] - 目标序列
            memory: [batch_size, src_seq_len, d_model] - 编码器输出
            src_mask: [batch_size, src_seq_len, src_seq_len] - 源序列掩码
            tgt_mask: [batch_size, tgt_seq_len, tgt_seq_len] - 目标序列掩码
        
        返回:
            output: [batch_size, tgt_seq_len, d_model] - 解码后的输出
        """
        # 掩码多头自注意力 + 残差连接 + 层归一化
        attn1_output, _ = self.self_attention(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout1(attn1_output))
        
        # 编码器-解码器注意力 + 残差连接 + 层归一化
        # 注意：交叉注意力中，查询来自目标序列，键值来自源序列
        # 所以掩码需要调整为 [batch_size, tgt_seq_len, src_seq_len]
        cross_attn_mask = None
        if src_mask is not None:
            # 源序列掩码的形状应该是 [batch_size, src_seq_len, src_seq_len]
            # 我们需要将其调整为 [batch_size, tgt_seq_len, src_seq_len]
            batch_size, tgt_seq_len = x.size(0), x.size(1)
            src_seq_len = memory.size(1)
            
            if src_mask.dim() == 3:
                # 对于3D掩码，我们只需要源序列的最后一个维度
                # 取掩码的最后一个维度，然后扩展以匹配目标序列长度
                cross_attn_mask = src_mask[:, -1:, :]  # [batch_size, 1, src_seq_len]
                cross_attn_mask = cross_attn_mask.expand(-1, tgt_seq_len, -1)  # [batch_size, tgt_seq_len, src_seq_len]
            else:
                # 对于2D掩码，直接扩展
                cross_attn_mask = src_mask.unsqueeze(0)  # [1, src_seq_len, src_seq_len]
                cross_attn_mask = cross_attn_mask.expand(batch_size, tgt_seq_len, -1)  # [batch_size, tgt_seq_len, src_seq_len]
        
        attn2_output, _ = self.cross_attention(x, memory, memory, cross_attn_mask)
        x = self.norm2(x + self.dropout2(attn2_output))
        
        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout3(ff_output))
        
        return x

"""解码器模块"""
class Decoder(nn.Module):
    """
    解码器
    由多个解码器层堆叠而成，包含位置编码和层归一化
    """
    
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, 
                 max_len=5000, dropout=0.1, **kwargs):
        super(Decoder, self).__init__(**kwargs)
        
        self.d_model = d_model
        
        # 词嵌入层
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # 位置编码
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 解码器层堆叠
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        # 输出层归一化
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        """
        前向传播
        
        参数:
            x: [batch_size, tgt_seq_len] - 目标序列的索引
            memory: [batch_size, src_seq_len, d_model] - 编码器输出
            src_mask: [batch_size, src_seq_len, src_seq_len] - 源序列掩码
            tgt_mask: [batch_size, tgt_seq_len, tgt_seq_len] - 目标序列掩码
        
        返回:
            output: [batch_size, tgt_seq_len, d_model] - 解码后的输出
        """
        # 词嵌入 + 缩放
        x = self.embedding(x) * math.sqrt(self.d_model)
        
        # 位置编码
        x = self.pos_encoding(x)
        
        # 通过所有解码器层
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        
        # 最终层归一化
        x = self.norm(x)
        
        return x

"""多头注意力机制模块"""
class MultiHeadAttention(nn.Module):
    """
    多头注意力机制
    将输入分割成多个头，在每个头上分别计算注意力，最后合并结果
    """
    
    def __init__(self, d_model, num_heads, dropout=0.1, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        
        # 参数验证
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model      # 输入维度
        self.num_heads = num_heads  # 注意力头数
        self.d_k = d_model // num_heads  # 每个头的维度
        
        # 线性变换层：将输入映射到Q、K、V
        self.W_q = nn.Linear(d_model, d_model)  # 查询变换
        self.W_k = nn.Linear(d_model, d_model)  # 键变换
        self.W_v = nn.Linear(d_model, d_model)  # 值变换
        
        # 输出线性变换层
        self.W_o = nn.Linear(d_model, d_model)
        
        # 缩放点积注意力模块
        self.attention = DotProductAttention(dropout)
        
    def forward(self, query, key, value, mask=None):
        """
            前向传播

            参数:
                query: [batch_size, seq_len_q, d_model]
                key: [batch_size, seq_len_k, d_model]
                value: [batch_size, seq_len_v, d_model]
                mask: [batch_size, seq_len_q, seq_len_k] 或 [seq_len_q, seq_len_k] - 可选掩码，用于屏蔽无效位置

            返回:
                output: [batch_size, seq_len_q, d_model]
                attention_weights: [batch_size, num_heads, seq_len_q, seq_len_k]
        """
        batch_size, seq_len_q, _ = query.size()
        _, seq_len_k, _ = key.size()
        
        # 步骤1: 线性变换得到Q、K、V
        Q = self.W_q(query)  # [batch_size, seq_len_q, d_model]
        K = self.W_k(key)    # [batch_size, seq_len_k, d_model]
        V = self.W_v(value)  # [batch_size, seq_len_v, d_model]
        
        # 步骤2: 分割成多个头
        # 重塑形状: [batch_size, seq_len, num_heads, d_k]
        # 转置: [batch_size, num_heads, seq_len, d_k]
        Q = Q.view(batch_size, seq_len_q, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)
        
        # 步骤3: 为每个头准备掩码（如果需要）
        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(0)  # 2D转3D [1, seq_len_q, seq_len_k]
            elif mask.dim() == 3 and mask.size(0) == 1:
                mask = mask.expand(batch_size, -1, -1)  # 扩展到批次
            # 扩展掩码以匹配多头形状
            mask = mask.unsqueeze(1)  # [batch_size, 1, seq_len_q, seq_len_k]
            mask = mask.expand(-1, self.num_heads, -1, -1)  # 扩展到每个头
        
        # 步骤4: 对每个头计算缩放点积注意力
        # 注意：这里我们使用批量矩阵乘法，同时处理所有头
        output, attention_weights = self.attention(Q, K, V, mask)
        
        # 步骤5: 合并多头输出
        # 转置回: [batch_size, seq_len_q, num_heads, d_k]
        output = output.transpose(1, 2).contiguous()
        # 重塑: [batch_size, seq_len_q, d_model]
        output = output.view(batch_size, seq_len_q, self.d_model)
        
        # 步骤6: 输出线性变换
        output = self.W_o(output)
        
        return output, attention_weights

"""完整的Transformer模型"""
class Transformer(nn.Module):
    """
    完整的Transformer模型
    包含编码器和解码器，用于序列到序列的任务
    """
    
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, 
                 d_ff=2048, num_layers=6, max_len=5000, dropout=0.1, **kwargs):
        super(Transformer, self).__init__(**kwargs)
        
        # 编码器
        self.encoder = Encoder(src_vocab_size, d_model, num_heads, d_ff, 
                               num_layers, max_len, dropout)
        
        # 解码器
        self.decoder = Decoder(tgt_vocab_size, d_model, num_heads, d_ff, 
                               num_layers, max_len, dropout)
        
        # 输出线性层（将解码器输出映射到目标词汇表大小）
        self.output_layer = nn.Linear(d_model, tgt_vocab_size)
        
        # 初始化参数
        self._init_weights()
    
    def _init_weights(self):
        """初始化模型参数"""
        # Xavier初始化线性层
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0, std=0.02)
    
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        """
        前向传播
        
        参数:
            src: [batch_size, src_seq_len] - 源序列索引
            tgt: [batch_size, tgt_seq_len] - 目标序列索引
            src_mask: [batch_size, src_seq_len, src_seq_len] - 源序列掩码
            tgt_mask: [batch_size, tgt_seq_len, tgt_seq_len] - 目标序列掩码
        
        返回:
            output: [batch_size, tgt_seq_len, tgt_vocab_size] - 预测概率分布
        """
        # 编码器前向传播
        memory = self.encoder(src, src_mask)
        
        # 解码器前向传播
        decoder_output = self.decoder(tgt, memory, src_mask, tgt_mask)
        
        # 输出层映射到词汇表大小
        output = self.output_layer(decoder_output)
        
        return output
    
    def generate_square_subsequent_mask(self, sz):
        """
        生成自回归掩码（下三角矩阵）
        用于解码器的自回归训练
        
        参数:
            sz: 序列长度
        
        返回:
            mask: [sz, sz] - 下三角掩码矩阵
        """
        mask = torch.triu(torch.ones(sz, sz), diagonal=1)
        mask = mask.masked_fill(mask == 1, float('-inf'))
        return mask


if __name__ == "__main__" :

    # 测试基本组件
    ffn = PositionWiseFFN(4, 4, 8)
    print("PositionWiseFFN测试:")
    print(ffn.eval())
    print(ffn(torch.ones((2, 3, 4))))
    print(ffn(torch.ones((2, 3, 4))).shape)
    
    # 测试多头注意力
    print("\nMultiHeadAttention测试:")
    mha = MultiHeadAttention(d_model=512, num_heads=8)
    query = torch.randn(2, 10, 512)
    key = torch.randn(2, 10, 512)
    value = torch.randn(2, 10, 512)
    output, weights = mha(query, key, value)
    print(f"输入形状: query{query.shape}, key{key.shape}, value{value.shape}")
    print(f"输出形状: {output.shape}, 注意力权重形状: {weights.shape}")
    
    # 测试编码器
    print("\nEncoder测试:")
    encoder = Encoder(vocab_size=1000, d_model=512, num_heads=8, 
                     d_ff=2048, num_layers=2)
    src = torch.randint(0, 1000, (2, 10))
    encoder_output = encoder(src)
    print(f"编码器输入形状: {src.shape}")
    print(f"编码器输出形状: {encoder_output.shape}")
    
    # 测试解码器
    print("\nDecoder测试:")
    decoder = Decoder(vocab_size=1000, d_model=512, num_heads=8, 
                     d_ff=2048, num_layers=2)
    tgt = torch.randint(0, 1000, (2, 8))
    memory = torch.randn(2, 10, 512)
    decoder_output = decoder(tgt, memory)
    print(f"解码器输入形状: tgt{tgt.shape}, memory{memory.shape}")
    print(f"解码器输出形状: {decoder_output.shape}")
    
    # 测试完整Transformer
    print("\nTransformer测试:")
    transformer = Transformer(src_vocab_size=1000, tgt_vocab_size=1000, 
                             d_model=512, num_heads=8, num_layers=2)
    src = torch.randint(0, 1000, (2, 10))
    tgt = torch.randint(0, 1000, (2, 8))
    
    # 生成目标序列掩码
    tgt_mask = transformer.generate_square_subsequent_mask(tgt.size(1))
    
    output = transformer(src, tgt, tgt_mask=tgt_mask)
    print(f"Transformer输入形状: src{src.shape}, tgt{tgt.shape}")
    print(f"Transformer输出形状: {output.shape}")
    print(f"目标序列掩码形状: {tgt_mask.shape}")
    
    print("\n所有组件测试完成！")